import rsa
import base64
from OpenSSL.crypto import PKey
from OpenSSL.crypto import TYPE_RSA, FILETYPE_PEM, FILETYPE_ASN1
from OpenSSL.crypto import dump_privatekey, dump_publickey

from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex


class PyRSA(object):

    @classmethod
    def gen_keys(cls):
        pk = PKey()
        pk.generate_key(TYPE_RSA, 2048)
        pub = dump_publickey(FILETYPE_PEM, pk)
        pri = dump_privatekey(FILETYPE_ASN1, pk)

        _pub_key = rsa.PublicKey.load_pkcs1_openssl_pem(pub)
        _private_key = rsa.PrivateKey.load_pkcs1(pri, 'DER')
        return _pub_key, _private_key, pub, pri

    @classmethod
    def parse_public_key(cls, pk):
        _pub_key = rsa.PublicKey.load_pkcs1_openssl_pem(pk)
        return _pub_key

    @classmethod
    def parse_private_key(cls, pk):
        _private_key = rsa.PrivateKey.load_pkcs1(pk, 'DER')
        return _private_key


class PyCrypto(object):
    def __init__(self, key):
        self.key = key
        self.mode = AES.MODE_CBC
        self.length = str.__len__(key)
        self.IV = 'iR9hPO63zXohlw2p'

    def encrypt(self, text):
        aes = AES.new(self.key, self.mode, self.IV)
        # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
        length = self.length
        count = len(text)

        add = 0
        if count % length != 0:
            add = length - (count % length)

        text = text + ('\0' * add)
        cipher_text = aes.encrypt(text)
        return base64.b64encode(cipher_text)

    def decrypt(self, text):
        aes = AES.new(self.key, self.mode, self.IV)

        # length = self.length
        # count = len(text)
        # add = 0
        # if count % length != 0:
        #     add = length - (count % length)
        #
        # text = text + ('\0' * add)

        plain_text = aes.decrypt(base64.b64decode(text))

        return plain_text.decode().strip('\0')


# (pubkey, prikey, pub, pri) = PyRSA.gen_keys()
#
# print(pubkey.save_pkcs1())
# print(prikey.save_pkcs1())

# data = rsa.encrypt(b'hello', pubkey)
# data = base64.b64encode(data)
#
# print(data)
#
# data0 = rsa.decrypt(base64.b64decode(data), prikey)
# print(data0)

# pem = b'-----BEGIN PUBLIC KEY-----\nMDwwDQYJKoZIhvcNAQEBBQADKwAwKAIhANJbJK1taeMcs4mGmaQxzouBbRCue9TG\n41ecfoDtv3DTAgMBAAE=\n-----END PUBLIC KEY-----\n'
# pem = '-----BEGIN PUBLIC KEY-----\nMDwwDQYJKoZIhvcNAQEBBQADKwAwKAIhAJWtQeGX+pVgChEANOWgf79pXGQgPPfORJ4tFTPCPxKnAgMBAAE=\n-----END PUBLIC KEY-----\n'
# publicKey = rsa.PublicKey.load_pkcs1_openssl_pem(pem)
# data = rsa.encrypt(b'hello', publicKey)
# data = base64.b64encode(data)

# print(data)
#
# data0 = rsa.decrypt(base64.b64decode(data), prikey)
# print(data0)




# pc = PyCrypto('fg5tred4qaedasds')
# d = pc.encrypt('http://www.baidu.com')
# print(d)
#
# dd = pc.decrypt('N25Lxjsa8Zq5aPA2xLtckxbF1n0HUqnXBWm8yaXjJ5k=')
# print(dd)
# print(len(dd))
